# A PyQt widget to display 3D scenes in OpenGL, along with interactive camera.
# Robert Bridson
# Released under the GPLv3

from PyQt4.QtCore import *
from PyQt4.QtGui import *
from PyQt4.QtOpenGL import *
from OpenGL.GL import *
import transform
import boundingbox
import math

class InteractiveCamera(object):

    # define some constants for different mouse mappings
    rotate = 1
    pan = 2
    dolly = 4

    def __init__(self):
        super(InteractiveCamera, self).__init__()
        self.fovy = 0.69813170079772 # 40 degrees in radians
        self.reset(1.0, boundingbox.bBox((-1,-1,-1), (1,1,1)))

    def reset(self, aspect, bbox):
        self.aspect = aspect
        self.bbox = bbox
        self.target, radius = boundingbox.sphere(bbox)
        self.distance = 3*radius
        self.near = 2*radius
        self.far = 4*radius
        self.heading = 0.0
        self.pitch = 0.0
        self._updateModelView()
        self.p = transform.perspective(self.fovy, self.aspect,
                                       self.near, self.far)

    def updateAspect(self, aspect):
        self.aspect = aspect
        self.p = transform.perspective(self.fovy, self.aspect,
                                       self.near, self.far)

    def updateSceneBoundingBox(self, bbox):
        if bbox != None:
            self.bbox = bbox
            # get vector to furthest corner from target
            d0 = max(self.target[0]-bbox[0][0], bbox[1][0]-self.target[0])
            d1 = max(self.target[1]-bbox[0][1], bbox[1][1]-self.target[1])
            d2 = max(self.target[2]-bbox[0][2], bbox[1][2]-self.target[2])
            # and the distance to the furthest corner from target
            r = math.sqrt(d0*d0+d1*d1+d2*d2)
            # then update near and far accordingly
            self.near = max(self.distance - r, 0.01*self.distance)
            self.far = self.distance + r

    def _updateModelView(self):
        t1 = transform.translate([-self.target[0],
                                  -self.target[1],
                                  -self.target[2]])
        t2 = transform.rotate_y(self.heading)
        t3 = transform.rotate_x(self.pitch)
        t4 = transform.translate([0, 0, -self.distance])
        self.mv = transform.compose(t4, t3, t2, t1)

    def mousePressEvent(self, event, w, h):
        self.mousePressPos = event.posF()
        mods = event.modifiers()
        if not bool(mods):
            self.dragMode = InteractiveCamera.rotate
            self.oldHeading = self.heading
            self.oldPitch = self.pitch
            self.mappingScale = math.pi/float(h)
        elif mods == Qt.ShiftModifier:
            self.dragMode = InteractiveCamera.pan
            self.oldTarget = self.target
            self.mappingScale = self.distance/float(h)
        elif mods == Qt.AltModifier:
            self.dragMode = InteractiveCamera.dolly
            self.oldDistance = self.distance
            self.mappingScale = 1.0/float(h) # appears in an exponent...
        else:
            self.dragMode = 0
        return False # nothing about the camera has changed yet...

    def mouseMoveEvent(self, event):
        diff = event.posF() - self.mousePressPos
        if self.dragMode == InteractiveCamera.rotate:
            self.heading = math.fmod(
                                  self.mappingScale*diff.x() + self.oldHeading,
                                  2*math.pi)
            self.pitch = min(math.pi/2, max(-math.pi/2,
                                    self.mappingScale*diff.y() + self.oldPitch))
            self._updateModelView()
            return True
        elif self.dragMode == InteractiveCamera.pan:
            # not implemented yet!
            return False
        elif self.dragMode == InteractiveCamera.dolly:
            # not implemented yet!
            return False
        else:
            return False # no need to trigger a re-render

    def mouseReleaseEvent(self, event):
        self.dragMode = 0
        return False # nothing about the camera has changed here

    def modelView(self):
        return self.mv

    def projection(self):
        return self.p


class View3d(QGLWidget):
    """
    View3d allows basic 3d rendering, taking care of some things for you.
    You can use the default interactive camera, or another camera object.
    Register objects to render, as long as they have initializeGL(),
    initialized(), boundingBox(), and paintGL(projViewTransform) methods
    (see axes.py for example), using addItem / removeItem.
    """

    def __init__(self, parent = None):
        QGLWidget.__init__(self, parent)
        self.camera = None
        self.aspectLocked = False # can vary as widget is resized
        self.aspect = 1.0 # save this to tell new cameras
        self.glStatus = False # so we know when we can do GL stuff
        self.items = set() # initially no items to render
        self.bbox = None # no knowledge yet of scene's bounding box

    def setCamera(self, camera):
        self.camera = camera
        self.camera.updateAspect(self.aspect)
        self.camera.updateSceneBoundingBox(self.bbox)
        self.update()

    def lockAspectRatio(self):
        self.aspectLocked = True

    def unlockAspectRatio(self):
        self.aspectLocked = False

    def addItem(self, item):
        self.items.add(item)

        if self.glStatus:
            item.initializeGL()
        if self.bbox == None:
            self.bbox = item.boundingBox()
        else:
            self.bbox = boundingbox.united(self.bbox, item.boundingBox())

        if self.camera != None:
            self.camera.updateSceneBoundingBox(self.bbox)

        self.update()

    def removeItem(self, item):
        self.items.discard(item)
        # recompute scene bounding box - not so efficient...
        self.bbox = None
        for x in self.items:
            if self.bbox == None:
                self.bbox = x.boundingBox()
            else:
                self.bbox = boundingbox.united(self.bbox, x.boundingBox())

    def mousePressEvent(self, event):
        if self.camera != None:
            if self.camera.mousePressEvent(event, self.w, self.h):
                self.update()
   
    def mouseMoveEvent(self, event):
        if self.camera != None and self.camera.mouseMoveEvent(event):
            self.update()

    def mouseReleaseEvent(self, event):
        if self.camera != None and self.camera.mouseReleaseEvent(event):
            self.update()

    def initializeGL(self):
        glClearDepth(1)
        for x in self.items:
            if not x.initialized():
                x.initializeGL()
        self.glStatus = True

    def resizeGL(self, w, h):
        glViewport(0, 0, w, h)
        if self.aspectLocked:
            # get ourselves a black background
            glClearColor(0, 0, 0, 0)
            glClear(GL_COLOR_BUFFER_BIT)
            # then figure out how to squeeze locked aspect ratio into widget
            widgetAspect = float(w)/float(h)
            if widgetAspect >= self.aspect: # widget too wide?
                lockedW = self.aspect*h
                margin = 0.5*(w-lockedW)
                glViewport(margin, 0, lockedW, h)
                self.w = lockedW
                self.h = h
            else: # widget too tall
                lockedH = w/self.aspect
                margin = 0.5*(h-lockedH)
                glViewport(0, margin, w, lockedH)
                self.w = w
                self.h = lockedH
        else:
            self.aspect = float(w)/float(h)
            self.w = w
            self.h = h
            # and get projection matrix to match
            if self.camera != None:
                self.camera.updateAspect(self.aspect)

    def paintGL(self):
        glClearColor(.75, .75, .75, 0)
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        if self.camera != None:
            pmv = transform.compose(self.camera.projection(),
                                    self.camera.modelView())
        else:
            pmv = transform.scale(1, self.aspect, 1)
        for x in self.items:
            x.paintGL(pmv)

